package com.qianheng.kit.base.util;

import cn.hutool.core.lang.Assert;
import com.qianheng.kit.base.dto.BizContext;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.ListUtils;
import org.springframework.lang.Nullable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

@Slf4j
public class AsyncThreadUtil {

    public static <R, Q> List<R> splitSupplyAsync(final List<Q> params,
        final int size,
        Executor executor,
        Function<List<Q>, List<R>> function,
        @Nullable BiConsumer<Exception, List<Q>> errorHandler) {

        if (CollectionUtils.isEmpty(params)) {
            return Collections.emptyList();
        }
        if (params.size() <= size) {
            return function.apply(params);
        }
        List<List<Q>> partition = ListUtils.partition(params, size);
        List<List<R>> supplyAsync = supplyAsync(partition, executor, function, errorHandler);
        return supplyAsync.stream().flatMap(Collection::stream).collect(Collectors.toList());
    }

    public static <R, Q> List<R> supplyAsync(final List<Q> params,
        Executor executor,
        Function<Q, R> function,
        @Nullable BiConsumer<Exception, Q> errorHandler) {
        if (CollectionUtils.isEmpty(params)) {
            return Collections.emptyList();
        }
        Assert.notNull(executor, "必须指定线程池");

        List<R> result = new ArrayList<>();
        List<QFu<Q, CompletableFuture<ThreadResHolder<R>>>> futures = new ArrayList<>();
        for (Q param : params) {
            CompletableFuture<ThreadResHolder<R>> f = CompletableFuture.supplyAsync(
                () -> new ThreadResHolder<>(function.apply(param), ObjUtil.deepCopy(TL.optBizContext().orElse(null), BizContext.class)),
                executor);
            futures.add(new QFu<>(param, f));
        }
        for (QFu<Q, CompletableFuture<ThreadResHolder<R>>> qfu : futures) {
            try {
                ThreadResHolder<R> rThreadResHolder = qfu.getFuture().get();
                result.add(rThreadResHolder.getRes());
                TL.addSubBizContextIfPresent(rThreadResHolder.getCtx());
            } catch (Exception e) {
                if (errorHandler != null) {
                    errorHandler.accept(e, qfu.getReq());
                } else {
                    log.error("splitSupplyAsync 异步执行异常", e);
                }
            }
        }
        return result;
    }

    @Getter
    @AllArgsConstructor
    public static class QFu<Q, Fu> {
        private Q req;
        private Fu future;
    }

    @Getter
    @AllArgsConstructor
    public static class ThreadResHolder<T> {
        private T res;
        private BizContext ctx;
    }
}
